from scapy import *
#from framework.capture__ import *
from data_type import * 
from datetime import *
import pickle as pickle
from string import *

def dosDetection(threshold):
    flowForDosD = []
    
    ###########################从pkl文件中获取flowForDosD列表###########################
    Ethernet = EthernetPkt()
    Ipv4 = Ipv4Pkt()
    Ipv6 = Ipv6Pkt()
    Tcp = TcpPkt()
    Udp = UdpPkt()
    Arp = ArpPkt()
    Dns = DnsPkt()
    Icmpv4 = Icmpv4Pkt()
    Icmpv6 = Icmpv6Pkt()
    HttpReq = HttpReqPkt()
    HttpRes = HttpResPkt()
    
    eachPkt = Packet(Ethernet, Ipv4, Ipv6, Tcp, Udp, Arp, Dns, Icmpv4, Icmpv6, HttpReq, HttpRes)
        
    #with open("C://Users//leslie//Desktop//res//cur//demo.pkl",'rb') as f:
    #with open("sniffRes//cur//demo.pkl",'rb') as f:
    with open("attack_demo//demo.pkl",'rb') as f:
        while 1:
            try:
                eachPkt = pickle.load(f)
                flowForDosD.append(eachPkt)
                #print(eachPkt)
                # eachPkt.PrintFlag()
            except EOFError:
                break
    ###########################从pkl文件中获取flowForDosD列表###########################
    
    thres = threshold
    packet_stream = {}      #该字段存储每个流以及对应的包数
    return_list = []
    count = 0
    for pkt in flowForDosD:
        if pkt.IsTcp():
            #pkt1=pkt.GetTcp()
            #print(pkt1)
            #if pkt.IsIpv4()
            #if type(pkt1)==Ipv4Pkt or type(pkt1)==Ipv6Pkt:
            if pkt.IsIpv4():
                src=pkt.GetIpv4().GetField('srcIp')
                dst=pkt.GetIpv4().GetField('dstIp')
                dport=pkt.GetTcp().GetField('dport')
            elif pkt.IsIpv6():
                src=pkt.GetIpv6().GetField('srcIp')
                dst=pkt.GetIpv6().GetField('dstIp')
                dport=pkt.GetTcp().GetField('dport')
            stream = str(src) + "-" + str(dst) + "-" +str(dport)
            #print(stream)
            if stream in  packet_stream.keys():
                packet_stream[stream] = packet_stream[stream] + 1
            else:
                packet_stream[stream] = 1 
            if len(packet_stream) > 0:
                for k, v in packet_stream.items():
                    if v > thres:
                        count += 1
                        print("DDOS attack found: \n%s times: %d" % (k, v))
                        k = k.split("-")
                        return_list.append((count, k[0], k[1], k[2], v))
                        # return k
                        #print(message)
                    #print(k,'\t', v)
    return return_list
    
if __name__ == "__main__":  
    print(dosDetection(2))
